/**
 * @file Multgrid_2D.h
 * @author your name (you@domain.com)
 * @brief 
 * @version 0.1
 * @date 2020-10-19
 * 
 * @copyright Copyright (c) 2020
 * 
 */
#ifndef _Multigrid_2D_H_
#define _Multigrid_2D_H_

#include <vector>
#include <algorithm>
#include <stdio.h>
#include <iostream>
#include <string>
#include <cstring>
#include <cmath>
#include <math.h>
#include "Multigrid.h"
#include "Miscellaneous.h"
#include "Restrict_2D.h"
#include "Interpolation_2D.h"

namespace MyMultigrid_2D
{
    /**
     * @brief The options of this Multigrid solver
     *       PROVIDEDOPTIONS:  
     *       1:BoundoryConditionType:Homogeneous(default) OR Nohomogeneous;
     *       2:RestrictionOperatorType:Full-weighting(default) OR Injection;
     *       3:InterpolationOperatorOperatorType:Linear(default) OR Quadratic;
     *       4:CycleType:Vcycle(default) OR FMcycle;
     *       5:MaxIterNum:Integer(default 10);
     *       6:RelAccurcy:double(default 1e - 6);
     */ 
    struct Options
    {
        std::string BoundoryConditionType = "Homogeneous";
        std::string RestrictionOperatorType = "Full-weighting";
        std::string InterpolationOperatorOperatorType = "Linear";
        std::string CycleType = "FMcycle";
        int MaxIterNum = 10;
        double RelAccurcy = 1e-5;
    };

    #define TEMPLATE template<int N>
    /**
     * @brief Template Multigrid solver for 2D problem
     *        N: The Number of segments in interval [0 , 1] × [0 , 1];
     * 
     */
    TEMPLATE
    class Multigrid_2D
    {
    private:
        /**
        * @brief Storage the variable of v;
        */
        std::vector<std::vector<double> > v;
        /**
         * @brief Storage the variable of f;
         */
        std::vector<std::vector<double> > f;
        /**
         * @brief torage the solution solved by Mutigrid method for each iteration;
         */
        std::vector<double> solution;
        /**
         * @brief Storage the  true solution of this equation;
         */
        std::vector<double> real_solution;
        /**
         * @brief Set the option of this Mutigrid method;
         */
        Options option;
        /**
         * @brief Storage the real solution function of this problem;
         *         Pointer to function;
         */
        double(*real)(double,double);
        /**
         * @brief Storage the righthands terms function of possion equation
         *        Pointer to function;
         */
        double(*rhf)(double,double);
        /**
         * @brief The highest level of the N;
         */
        int level = int(log(N) / log(2));
        /**
         * @brief current relativeaccuracy;
         */
        double relativeaccuracy;
        /**
         * @brief current iteration number
         */
        int IterNumber = 0;
        /**
         * @brief Selection of restrict operator
         */
        MyMultigrid_2D::Operator::RestrictOperator_2D * restrict;
        /**
         * @brief Selection of Interpolation operator
         */
        MyMultigrid_2D::Operator::InterpolationOperator_2D *interpolation;
        
    public:
        /**
         * @brief Construct a new Multigrid_2D object
         */
        Multigrid_2D(){};
        /**
         * @brief Construct a new Multigrid_2D object
         * 
         * @param u u real soution function of possion equation
         * @param func righthands term function of possion equation
         * @param opt The options of Multigrid method
         */
        Multigrid_2D(double(*u)(double,double),double(*func)(double,double),MyMultigrid_2D::Options opt);
        /**
         * @brief Destroy the Multigrid_2D object
         */
        ~Multigrid_2D(){};
        /**
         * @brief The Weighted-jacobi iteration method
         * 
         * @param l current level
         * @param times iteration times
         * @param omage weight
         */
        void solve();

        template<int N_>
        friend std::ostream& operator<<(std::ostream& os,MyMultigrid_2D::Multigrid_2D<N_> m);  

        private:
        void weightedjacobi(int l,int times,double omage,int s);
        void Vcycle(int l,int s);
        void FMcycle(int l);
        void reset_v(int l){for(int k = 1;k < l;k ++) v[k].assign(v[k].size(),0);};
        void reset_f(int l){for(int k = 1;k < l;k ++) f[k].assign(v[k].size(),0);};
        
    };
    TEMPLATE
    Multigrid_2D<N>::Multigrid_2D(double(*u)(double,double),double(*func)(double,double),Options opt)
    {
        v.resize(level + 1);
        f.resize(level + 1);
        solution.resize((pow(2,level)+1)*(pow(2,level)+1) );
        real_solution.resize((pow(2,level)+1)*(pow(2,level)+1));
        for(int k = 1;k <=level;k++)
        {
            v[k].resize((pow(2,k)+1)*(pow(2,k)+1));
            f[k].resize((pow(2,k)+1)*(pow(2,k)+1));
        }
        for(int i = 0;i < f[level].size();i++)
        {
            double h = 1 / (pow(2,level));
            f[level][i] = func( index_i(i,level)* h,index_j(i,level) * h);
            real_solution[i] = u(index_i(i,level)* h,index_j(i,level) * h);
        }
        real = u;
        rhf = func;
        option = opt;
        restrict = new MyMultigrid_2D::Operator::FullWeightingRestrictOperator_2D();
        interpolation = new MyMultigrid_2D::Operator::LinearInterpolaionOperator_2D();
        
    };

    TEMPLATE
    void Multigrid_2D<N>::Vcycle(int l ,int s)
    {
        weightedjacobi(l,2,2/3.0,s);
        if(l > 1)
        {
            double h_ = 1 / double(pow(2,l));
            std::vector<double> r_h(v[l].size());
            for(int k = 0;k < v[l].size();k++)
            {
                if(isboundary(k,l) == false)
                {
                    int i = index_i(k,l);
                    int j = index_j(k,l);
                    int v1 = index(i - 1,j,l);
                    int v2 = index(i + 1,j,l);
                    int v3 = index(i,j + 1,l);
                    int v4 = index(i,j - 1,l);
                    int v5 = index(i,j,l);
                    r_h[k] =(f[l][k] + ((v[l][v1] + v[l][v2] + v[l][v3] + v[l][v4] - 4 * v[l][v5]) /(h_*h_)));
                }
            }
            restrict->SetInput(r_h);
            restrict ->DoRestrict(l);
            std::vector<double> r_2h = restrict ->ReturnRestricResult();
            f[l - 1] = r_2h;
            Vcycle(l - 1,s);
        }
        else
        {
            double h = 1.0 /(double(pow(2,l)));
            v[l][4] = 0.25*(h * h * f[l][4]);
        }

        if(l != s)
        {
            std::vector<double> r_2h_ = v[l];
            interpolation ->SetInput(r_2h_);
            interpolation ->DoInterpolation(l);
            std::vector<double> r_h_ = interpolation -> ReturnInterpolaionResult();
            for(int i = 0;i < v[l + 1].size();i++)
            {
                v[l+1][i] += r_h_[i];
            }
            weightedjacobi(l+1,2,2/3.0,s);
        }
        
    }

    TEMPLATE
    void Multigrid_2D<N>::FMcycle(int l)
    {
        if(l > 1)
        {
            restrict -> SetInput(f[l]);
            restrict -> DoRestrict(l);
            f[l - 1] =  restrict -> ReturnRestricResult();
            FMcycle(l - 1);
        }
        else
        {
            double h = 1.0 /(double(pow(2,l)));
            v[l][4] = 0.25*(h * h * f[l][4] +  v[l][1]  + v[l][3] + v[l][5]  + v[l][7]);
        }
        if(l != level)
        {
            interpolation -> SetInput(v[l]);
            interpolation -> DoInterpolation(l);
            v[l + 1] = interpolation -> ReturnInterpolaionResult();
            for(int i = 0;i < 6;i ++)
            {
                reset_v(l + 1);
                Vcycle(l + 1,l + 1);
            }
        }
        
    };
    TEMPLATE
    void Multigrid_2D<N>::solve()
    {
        double h = 1.0 / double(pow(2,level));
        std::vector<double> e(real_solution.size());
        for(int i = 0;i < e.size();i++)
        {
            e[i] = real_solution[i] - solution[i];
        }
        double acctemp = norm_max(e);
        int currenttimes = 0;
        std::cout << "..Calculating.."<<std::endl;
        do
        {
            if(option.CycleType == "Vcycle")
                Vcycle(level,level);
            else
                FMcycle(level);
            solution = v[level];
            for(int i = 0;i < e.size();i++)
            {
               e[i] = real_solution[i] - solution[i];
            }
            acctemp = norm_max(e);
            currenttimes++;
            reset_f(level);
            reset_v(level);            
            std::cout <<"Current iteration time is :" << currenttimes << std::endl;
            printf("Current relative accurcy is : %.8lf\n",acctemp);
        }while (currenttimes < option.MaxIterNum && acctemp > option.RelAccurcy);
        relativeaccuracy = acctemp;
        IterNumber = currenttimes;
        if(currenttimes >= option.MaxIterNum && relativeaccuracy > option.RelAccurcy)
        {
            std::cout << "Iteration number arrive the Max Number!"<<std::endl;
            std::cout << "But does not convergence to the preset relative accrucy!"<<std::endl;
        }
        else
        {
            std::cout << "Iteration convergence to " << option.RelAccurcy << " at "<< IterNumber << "th time"<<std::endl;
            printf("The last relative accurcy is : %.8lf !\n",relativeaccuracy);
        } 
        std::cout << "..Done!.."<<std::endl;
    }; 

    TEMPLATE
    void Multigrid_2D<N>::weightedjacobi(int l,int times,double omage,int s)
    {
        std::vector<double> temp;
        if(l == s)
        {
            double h = 1 / (pow(2,l));
            for(int k = 0;k < v[l].size();k++)
            {
                if(isboundary(k,l))
                {

                    int i = index_i(k,l);
                    int j = index_j(k,l);
                    v[l][k] = real(i * h,j * h);
                }
            }
        }
        else
        {
            for(int i = 0;i < v[l].size();i++)
            {
                if(isboundary(i,l))
                    v[l][i] = 0;
            }
        }
        double h = 1 / double(pow(2,l));
        for(int k = 0;k < times;k++)
        {
            temp = v[l];
            for(int i = 0;i < v[l].size();i++)
            {
                if(isboundary(i,l) == false)
                {
                    int a = index_i(i,l);
                    int b = index_j(i,l);
                    int c1 = index(a - 1,b,l);
                    int c2 = index(a + 1,b,l);
                    int c3 = index(a,b - 1,l);
                    int c4 = index(a,b + 1,l);
                    double t = 0.25 *(temp[c1] + temp[c2] + temp[c3] + temp[c4] + h * h * f[l][i]);
                    v[l][i] = temp[i] + omage * (t - temp[i]);
                }
            }
        }   
    }
    template<int N_>
    std::ostream& operator<<(std::ostream& os,::MyMultigrid_2D::Multigrid_2D<N_> m)
    {
        os << "MyMultigrid_2D::Multigrid_2D Options setting information !" <<std::endl;
        os << "The segments of the Multigrid solver : " << N_ <<"*" << N_ << std::endl;
        os << "The Options of the mutigrid solver : " << m.option.BoundoryConditionType <<std::endl;
        os << "                                     " <<m.option.RestrictionOperatorType << std::endl;
        os << "                                     " << m.option.InterpolationOperatorOperatorType << std::endl;
        os << "                                     " << m.option.CycleType<<std::endl;
        os << "                                     " << m.option.MaxIterNum<<std::endl;
        os << "                                     " <<m.option.RelAccurcy <<std::endl;
        return os;
    }
} // namespace MyMultigrid_2D

#else
#endif